package com.bruce.ai.alibaba.agentcard.discovery.auth;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.time.Instant;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * OAuth2令牌管理器
 * 负责OAuth2令牌的获取、刷新和缓存管理
 * 
 * 特性：
 * - 客户端凭证流程
 * - 令牌缓存和自动刷新
 * - 多客户端支持
 * - 线程安全
 */
@Component
public class OAuth2TokenManager {
    
    private static final Logger logger = LoggerFactory.getLogger(OAuth2TokenManager.class);
    private static final String GRANT_TYPE_CLIENT_CREDENTIALS = "client_credentials";
    
    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper;
    private final ConcurrentMap<String, TokenInfo> tokenCache;
    
    public OAuth2TokenManager() {
        this.httpClient = new OkHttpClient.Builder()
                .connectTimeout(10, java.util.concurrent.TimeUnit.SECONDS)
                .readTimeout(30, java.util.concurrent.TimeUnit.SECONDS)
                .build();
        this.objectMapper = new ObjectMapper();
        this.tokenCache = new ConcurrentHashMap<>();
    }
    
    /**
     * 获取访问令牌（客户端凭证流程）
     * @param tokenEndpoint 令牌端点URL
     * @param clientId 客户端ID
     * @param clientSecret 客户端密钥
     * @param scope 权限范围（可选）
     * @return 访问令牌
     * @throws OAuth2Exception OAuth2异常
     */
    public String getAccessToken(String tokenEndpoint, String clientId, String clientSecret, String scope) 
            throws OAuth2Exception {
        
        String cacheKey = generateCacheKey(tokenEndpoint, clientId, scope);
        TokenInfo tokenInfo = tokenCache.get(cacheKey);
        
        // 检查缓存的令牌是否有效
        if (tokenInfo != null && !tokenInfo.isExpired()) {
            logger.debug("Using cached access token for client: {}", clientId);
            return tokenInfo.getAccessToken();
        }
        
        // 获取新令牌
        logger.info("Requesting new access token for client: {}", clientId);
        tokenInfo = requestAccessToken(tokenEndpoint, clientId, clientSecret, scope);
        
        // 缓存令牌
        tokenCache.put(cacheKey, tokenInfo);
        
        return tokenInfo.getAccessToken();
    }
    
    /**
     * 刷新访问令牌
     * @param tokenEndpoint 令牌端点URL
     * @param clientId 客户端ID
     * @param clientSecret 客户端密钥
     * @param refreshToken 刷新令牌
     * @return 新的访问令牌
     * @throws OAuth2Exception OAuth2异常
     */
    public String refreshAccessToken(String tokenEndpoint, String clientId, String clientSecret, String refreshToken) 
            throws OAuth2Exception {
        
        logger.info("Refreshing access token for client: {}", clientId);
        
        FormBody.Builder formBuilder = new FormBody.Builder()
                .add("grant_type", "refresh_token")
                .add("refresh_token", refreshToken);
        
        Request request = new Request.Builder()
                .url(tokenEndpoint)
                .post(formBuilder.build())
                .header("Authorization", Credentials.basic(clientId, clientSecret))
                .header("Content-Type", "application/x-www-form-urlencoded")
                .build();
        
        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new OAuth2Exception("Failed to refresh token: " + response.code() + " " + response.message());
            }
            
            String responseBody = response.body().string();
            JsonNode jsonNode = objectMapper.readTree(responseBody);
            
            if (jsonNode.has("error")) {
                throw new OAuth2Exception("OAuth2 error: " + jsonNode.get("error").asText());
            }
            
            String accessToken = jsonNode.get("access_token").asText();
            int expiresIn = jsonNode.has("expires_in") ? jsonNode.get("expires_in").asInt() : 3600;
            
            // 更新缓存
            String cacheKey = generateCacheKey(tokenEndpoint, clientId, null);
            TokenInfo tokenInfo = new TokenInfo(accessToken, refreshToken, expiresIn);
            tokenCache.put(cacheKey, tokenInfo);
            
            return accessToken;
            
        } catch (IOException e) {
            throw new OAuth2Exception("Network error during token refresh", e);
        }
    }
    
    /**
     * 撤销访问令牌
     * @param revokeEndpoint 撤销端点URL
     * @param token 要撤销的令牌
     * @param clientId 客户端ID
     * @param clientSecret 客户端密钥
     * @throws OAuth2Exception OAuth2异常
     */
    public void revokeToken(String revokeEndpoint, String token, String clientId, String clientSecret) 
            throws OAuth2Exception {
        
        logger.info("Revoking token for client: {}", clientId);
        
        FormBody formBody = new FormBody.Builder()
                .add("token", token)
                .build();
        
        Request request = new Request.Builder()
                .url(revokeEndpoint)
                .post(formBody)
                .header("Authorization", Credentials.basic(clientId, clientSecret))
                .header("Content-Type", "application/x-www-form-urlencoded")
                .build();
        
        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new OAuth2Exception("Failed to revoke token: " + response.code() + " " + response.message());
            }
        } catch (IOException e) {
            throw new OAuth2Exception("Network error during token revocation", e);
        }
    }
    
    /**
     * 清理过期的令牌缓存
     */
    public void cleanupExpiredTokens() {
        int removedCount = 0;
        for (String key : tokenCache.keySet()) {
            TokenInfo tokenInfo = tokenCache.get(key);
            if (tokenInfo != null && tokenInfo.isExpired()) {
                tokenCache.remove(key);
                removedCount++;
            }
        }
        if (removedCount > 0) {
            logger.info("Cleaned up {} expired tokens from cache", removedCount);
        }
    }
    
    /**
     * 获取缓存统计信息
     * @return 缓存大小
     */
    public int getCacheSize() {
        return tokenCache.size();
    }
    
    /**
     * 清空令牌缓存
     */
    public void clearCache() {
        tokenCache.clear();
        logger.info("OAuth2 token cache cleared");
    }
    
    /**
     * 请求访问令牌
     */
    private TokenInfo requestAccessToken(String tokenEndpoint, String clientId, String clientSecret, String scope) 
            throws OAuth2Exception {
        
        FormBody.Builder formBuilder = new FormBody.Builder()
                .add("grant_type", GRANT_TYPE_CLIENT_CREDENTIALS);
        
        if (scope != null && !scope.trim().isEmpty()) {
            formBuilder.add("scope", scope);
        }
        
        Request request = new Request.Builder()
                .url(tokenEndpoint)
                .post(formBuilder.build())
                .header("Authorization", Credentials.basic(clientId, clientSecret))
                .header("Content-Type", "application/x-www-form-urlencoded")
                .build();
        
        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new OAuth2Exception("Failed to get access token: " + response.code() + " " + response.message());
            }
            
            String responseBody = response.body().string();
            JsonNode jsonNode = objectMapper.readTree(responseBody);
            
            if (jsonNode.has("error")) {
                throw new OAuth2Exception("OAuth2 error: " + jsonNode.get("error").asText());
            }
            
            String accessToken = jsonNode.get("access_token").asText();
            String refreshToken = jsonNode.has("refresh_token") ? jsonNode.get("refresh_token").asText() : null;
            int expiresIn = jsonNode.has("expires_in") ? jsonNode.get("expires_in").asInt() : 3600;
            
            return new TokenInfo(accessToken, refreshToken, expiresIn);
            
        } catch (IOException e) {
            throw new OAuth2Exception("Network error during token request", e);
        }
    }
    
    /**
     * 生成缓存键
     */
    private String generateCacheKey(String tokenEndpoint, String clientId, String scope) {
        return String.format("%s:%s:%s", tokenEndpoint, clientId, scope != null ? scope : "");
    }
    
    /**
     * 令牌信息类
     */
    private static class TokenInfo {
        private final String accessToken;
        private final String refreshToken;
        private final long expiresAt;
        
        public TokenInfo(String accessToken, String refreshToken, int expiresIn) {
            this.accessToken = accessToken;
            this.refreshToken = refreshToken;
            // 提前30秒过期以避免边界情况
            this.expiresAt = Instant.now().getEpochSecond() + expiresIn - 30;
        }
        
        public String getAccessToken() {
            return accessToken;
        }
        
        public String getRefreshToken() {
            return refreshToken;
        }
        
        public boolean isExpired() {
            return Instant.now().getEpochSecond() >= expiresAt;
        }
    }
    
    /**
     * OAuth2异常类
     */
    public static class OAuth2Exception extends Exception {
        public OAuth2Exception(String message) {
            super(message);
        }
        
        public OAuth2Exception(String message, Throwable cause) {
            super(message, cause);
        }
    }
}